package com.xiaozhou.controller;

import com.xiaozhou.function.FunctionManager;
import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.service.v4.tools.SearchChatMessage;
import com.zhipu.oapi.service.v4.tools.WebSearchApiResponse;
import com.zhipu.oapi.service.v4.tools.WebSearchParamsRequest;
import com.zhipu.oapi.service.v4.tools.WebSearchPro;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;

/**
 * @author shhzhou
 * @description: 专业联网搜索工具 范例
 * @date 2024/11/4 16:51
 */
@RestController
@RequestMapping("/webSearch")
public class WebSearchController {

    @Autowired
    private ClientV4 clientV4;

    @Autowired
    private FunctionManager functionManager;

    /**
     * @return
     * @return com.zhipu.oapi.service.v4.tools.WebSearchApiResponse
     * @Author xiaozhou
     * @Description 联网搜索 同步调用 例子
     * @Date 17:04 2024/11/4
     * @Param
     */
    @GetMapping("/syncModel")
    public WebSearchApiResponse syncModel(String input) {

        List<SearchChatMessage> messages = new ArrayList<>();
        SearchChatMessage message = new SearchChatMessage();
        message.setRole("user");
        message.setContent(input);

        messages.add(message);

        // 构造模型对象
        WebSearchParamsRequest request = new WebSearchParamsRequest();
        request.setModel("web-search-pro");
        request.setStream(false);
        request.setMessages(messages);


        WebSearchApiResponse webSearchApiResponse = clientV4.invokeWebSearchPro(request);
        return webSearchApiResponse;
    }


    /**
     * @return
     * @return reactor.core.publisher.Flux<com.zhipu.oapi.service.v4.tools.WebSearchPro>
     * @Author xiaozhou
     * @Description 联网搜索 sse流式调用 例子
     * @Date 17:08 2024/11/4
     * @Param
     */
    @GetMapping("/sseModel")
    public Flux<WebSearchPro> sseModel(String input) {

        List<SearchChatMessage> messages = new ArrayList<>();
        SearchChatMessage message = new SearchChatMessage();
        message.setRole("user");
        message.setContent(input);

        messages.add(message);

        // 构造模型对象
        WebSearchParamsRequest request = new WebSearchParamsRequest();
        request.setModel("web-search-pro");
        request.setStream(true);
        request.setMessages(messages);

        WebSearchApiResponse webSearchApiResponse = clientV4.webSearchProStreamingInvoke(request);

        // 使用 Flux.create 创建流
        return Flux.create(sink -> {
            webSearchApiResponse.getFlowable().map(modelData -> {
                return modelData;
            }).doOnNext(modelData -> {
                sink.next(modelData);
            }).doOnComplete(() -> {
                sink.complete();
            }).blockingSubscribe();
        });
    }


}


